import numpy as np
import torch

from torchvision.transforms import Compose
import torch.nn as nn
from utils.utils import Rescale, Normailize, Reshape
# from utils.nms_wrapper import nms
import time
import torch.nn.functional as F
import torch
import os
import numpy as np
import math
# from datasets.DOTA_devkit.ResultMerge_multi_process import py_cpu_nms_poly_fast, py_cpu_nms_poly
from PIL import Image, ImageDraw
import cv2
import sys
sys.path.append("..")
from loss.Loss import _gather_feat,_transpose_and_gather_feat
from datasets_loader.ssdd import get_affine_transform


def decode_prediction(img, res):
    ptss = []
    for class_name,lx,ly,rx,ry,ang, prob in res:

        result = [int((rx+lx)/2),int((ry+ly)/2),int(rx-lx),int(ry-ly),ang]
        result=np.array(result)
        x=int(result[0])
        y=int(result[1])
        height=int(result[3])
        width=int(result[2])

        anglePi = (result[4])/180 * math.pi

        anglePi = anglePi if anglePi <= math.pi else anglePi - math.pi
 
        cosA = math.cos(anglePi)
        sinA = math.sin(anglePi)
        
        x1=x-0.5*width   
        y1=y-0.5*height
        
        x0=x+0.5*width 
        y0=y1
        
        x2=x1            
        y2=y+0.5*height 
        
        x3=x0   
        y3=y2
        
        x0n= (x0 -x)*cosA -(y0 - y)*sinA + x
        y0n = (x0-x)*sinA + (y0 - y)*cosA + y
        
        x1n= (x1 -x)*cosA -(y1 - y)*sinA + x
        y1n = (x1-x)*sinA + (y1 - y)*cosA + y
        
        x2n= (x2 -x)*cosA -(y2 - y)*sinA + x
        y2n = (x2-x)*sinA + (y2 - y)*cosA + y
        
        x3n= (x3 -x)*cosA -(y3 - y)*sinA + x
        y3n = (x3-x)*sinA + (y3 - y)*cosA + y

        
        tr = np.asarray([x0n,y0n], np.float32)
        br = np.asarray([x3n,y3n], np.float32)
        bl = np.asarray([x2n,y2n], np.float32)
        tl = np.asarray([x1n,y1n], np.float32)
        cen_pt = np.asarray([x, y], np.float32)
        
        score = prob            
        clse = int(class_name)-1

        pts = [clse,score,x1n,y1n,x0n,y0n,x3n,y3n,x2n,y2n]     
        ptss.append(pts)

    return ptss

def pre_process(args,image):

   # cv2.imwrite('img_predict_results/'+"init.jpg",image)
    height, width = image.shape[0:2]
   # inp_height, inp_width = args.input_h,args.input_w
    c = np.array([height / 2., width / 2.], dtype=np.float32)
    
    if args.keep_res:
        inp_height = (height | 31) + 1   
        inp_width = (width | 31) + 1
#         print(input_h)
        s = np.array([inp_width, inp_height], dtype=np.float32)
    else:
        s = max(height,  width) * 1.0  
        inp_height, inp_width = args.input_h,args.input_w
        
#     s = max(height, width) * 1
#     s = np.array([height, width])
    trans_input = get_affine_transform(c, s, 0, [inp_width, inp_height])
    
 #将仿射矩阵最后一列置为0， 最后一列表示平移。  现仅有缩放
#####################################################################
    trans_input[:,-1]=0.
#####################################################################

    inp_image = cv2.warpAffine(image, trans_input, (inp_width, inp_height),flags=cv2.INTER_LINEAR)

#     cv2.imwrite('img_predict_results/'+"test.jpg",inp_image)

    mean = np.array([0.5194416012442385,0.5378052387430711,0.533462090585746], dtype=np.float32).reshape(1, 1, 3)
    std  = np.array([0.3001546018824507, 0.28620901391179554, 0.3014112676161966], dtype=np.float32).reshape(1, 1, 3)
#     mean = np.array([0.07005435, 0.07005435, 0.07005435], dtype=np.float32).reshape(1, 1, 3)
#     std  = np.array([0.05547728, 0.05547728, 0.05547728], dtype=np.float32).reshape(1, 1, 3)
    
    inp_image = ((inp_image / 255. - mean) / std).astype(np.float32)

    images = inp_image.transpose(2, 0, 1).reshape(1, 3, inp_height, inp_width) # 三维reshape到4维，（1，3，512，512） 
    
    images = torch.from_numpy(images)

    meta = {'c': c, 's': s, 
            'out_height': inp_height // 4, #取整除
            'out_width': inp_width // 4}
    return images, meta


def _nms(heat, kernel=3):
    pad = (kernel - 1) // 2
    hmax = nn.functional.max_pool2d(
        heat, (kernel, kernel), stride=1, padding=pad)

    keep = (hmax == heat).float()
    return heat * keep


def _topk(scores, K=40):
    batch, cat, height, width = scores.size()
    topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K)#取每一类的最多K个
    topk_inds = topk_inds % (height * width)
    
    topk_ys   = (topk_inds / width).int().float()#     topk_ys = (torch.true_divide(topk_inds, width)).int().float()
    
    topk_xs   = (topk_inds % width).int().float() 
    topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)
    
    topk_clses = (topk_ind / K).int()#     topk_clses = (torch.true_divide(topk_ind,K)).int()


    topk_inds = _gather_feat(
        topk_inds.view(batch, -1, 1), topk_ind).view(batch, K)
    topk_ys = _gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, K)
    topk_xs = _gather_feat(topk_xs.view(batch, -1, 1), topk_ind).view(batch, K)

    return topk_score, topk_inds, topk_clses, topk_ys, topk_xs


def ctdet_decode(args,heat, wh, ang,num_classes ,reg=None, K=100):
    batch, cat, height, width = heat.size()

    # heat = torch.sigmoid(heat)
    # perform nms on heatmaps
    heat = _nms(heat)
    scores, inds, clses, ys, xs = _topk(heat, K=K)

    reg = _transpose_and_gather_feat(reg, inds)
    reg = reg.view(batch, K, 2)
    xs = xs.view(batch, K, 1) + reg[:, :, 0:1]
    ys = ys.view(batch, K, 1) + reg[:, :, 1:2]
   
    wh = _transpose_and_gather_feat(wh, inds)
    wh = wh.view(batch, K, 2)
    
    ang = _transpose_and_gather_feat(ang, inds)

    ang = ang.view(batch, K, 1)


    clses  = clses.view(batch, K, 1).float()

    scores = scores.view(batch, K, 1)
    bboxes = torch.cat([xs - wh[..., 0:1] / 2, 
                        ys - wh[..., 1:2] / 2,
                        xs + wh[..., 0:1] / 2, 
                        ys + wh[..., 1:2] / 2,
                        ang], dim=2)
    detections = torch.cat([bboxes, scores, clses], dim=2)

    return detections


def process(args,images,num_classes, model,return_time=False):
    with torch.no_grad():
        output = model(images)
        hm = output['hm'].sigmoid_()
        ang = output['ang'].relu_()
        wh = output['wh']
        reg = output['reg'] 
        torch.cuda.synchronize()
        forward_time = time.time()
        dets = ctdet_decode(args,hm, wh, ang, num_classes,reg=reg, K=100) # K 是最多保留几个目标
    if return_time:
        return output, dets, forward_time
    else:
        return output, dets


def affine_transform(pt, t):
    new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32).T
    new_pt = np.dot(t, new_pt)
    return new_pt[:2]


def transform_preds(coords, center, scale, output_size):
    target_coords = np.zeros(coords.shape)
    trans = get_affine_transform(center, scale, 0, output_size, inv=1)
    
    
 #将仿射矩阵最后一列置为0， 最后一列表示平移。  现仅有缩放
#####################################################################
    trans[:,-1]=0.
#####################################################################    

    for p in range(coords.shape[0]):
        target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
    return target_coords


def ctdet_post_process(dets, c, s, h, w, num_classes):
  # dets: batch x max_dets x dim
  # return 1-based class det dict
    ret = []
#     print("dets.size()",dets.shape)#(1, 100, 7)
    for i in range(dets.shape[0]):
        top_preds = {}
        #print("dets[i, :, :2]:",dets[i, :, :2])
        dets[i, :, :2] = transform_preds(dets[i, :, 0:2], c[i], s[i], (w, h))
        #print("dets[i, :, :2]:",dets[i, :, :2])
        dets[i, :, 2:4] = transform_preds(dets[i, :, 2:4], c[i], s[i], (w, h))
        classes = dets[i, :, -1]
    for j in range(num_classes):
        inds = (classes == j)
      
        top_preds[j + 1] = np.concatenate([
        dets[i, inds, :4].astype(np.float32),
        dets[i, inds, 4:6].astype(np.float32)], axis=1).tolist()
#         print(top_preds)
        ret.append(top_preds)
    return ret


def post_process(dets, meta,num_classes):  
    dets = dets.detach().cpu().numpy()
#     print("dets.shape变化前:",dets.shape)
    dets = dets.reshape(1, -1, dets.shape[2])  
#     print("dets.shape变化后:",dets.shape)  
    
#     print("meta",meta)
    dets = ctdet_post_process(dets.copy(), [meta['c']], [meta['s']],meta['out_height'], meta['out_width'], num_classes)
    for j in range(1, num_classes + 1):
        dets[0][j] = np.array(dets[0][j], dtype=np.float32).reshape(-1, 6)
        dets[0][j][:, :5] /= 1
    return dets[0]


def merge_outputs(args,detections,num_classes):

    max_obj_per_img = args.max_obj_per_img

    scores = np.hstack([detections[j][:, 5] for j in range(1, num_classes + 1)])#在水平方向上平铺

    if len(scores) > max_obj_per_img:
        kth = len(scores) - max_obj_per_img
        thresh = np.partition(scores, kth)[kth]
        
        for j in range(1, 2 + 1):
            keep_inds = (detections[j][:, 5] >= thresh)        

            detections[j] = detections[j][keep_inds]
    return detections
    
    
    
    

def im_detect(args,image_name,model,down_ratio):
    num_classes =args.num_classes
    device = torch.device('cuda')
    image_ = cv2.imread(image_name)
    images, meta = pre_process(args,image_)#仿射变换---缩？放？
    images = images.to(device)
    output, dets, forward_time = process(args,images,num_classes=num_classes, return_time=True,model = model)
            
    dets = post_process(dets, meta,num_classes)#仿射变换，缩？放？
    ret = merge_outputs(args,dets,num_classes)
#     print(ret)
    res = np.empty([1,7])
    for i, c in ret.items():
        tmp_s = ret[i][ret[i][:,5]>0.1]
        tmp_c = np.ones(len(tmp_s)) * (i+1)
        tmp = np.c_[tmp_c,tmp_s]
        res = np.append(res,tmp,axis=0)

    res = np.delete(res, 0, 0)
    res = res.tolist()#tolist()作用：将矩阵（matrix）和数组（array）转化为列表。
    
    ptss = np.array(decode_prediction(image_, res))
    return ptss
#     print(res)
#     print(image_)
    
    
        






